#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>

// 确保矩阵维度是4的倍数（简化对齐问题）
#define M 16
#define K 16
#define N 16

// 验证结果是否正确的阈值
#define EPSILON 1e-5

// CPU参考计算（普通逐元素计算）
void matrixMulCPU(float *A, float *B, float *C) {
    for (int i = 0; i < M; ++i) {
        for (int j = 0; j < N; ++j) {
            float sum = 0.0f;
            for (int k = 0; k < K; ++k) {
                sum += A[i * K + k] * B[k * N + j];
            }
            C[i * N + j] = sum;
        }
    }
}

// GPU核函数（使用float4优化）
__global__ void matrixMulGPU(float4 *A, float4 *B, float *C) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float sum = 0.0f;

        // 每个float4包含4个元素，因此循环次数为K/4
        for (int k = 0; k < K/4; ++k) {
            // 从A加载一个float4（行方向）
            float4 a = A[row * (K/4) + k];
            // 从B加载一个float4（列方向，需要转置访问）
            float4 b = B[col * (K/4) + k];

            // 计算点积：a.x*b.x + a.y*b.y + a.z*b.z + a.w*b.w
            sum += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
        }

        C[row * N + col] = sum;
    }
}

int main() {
    // 主机内存分配（注意：必须是连续的4的倍数！）
    float *h_A = (float*)malloc(M * K * sizeof(float));
    float *h_B = (float*)malloc(K * N * sizeof(float));
    float *h_C = (float*)malloc(M * N * sizeof(float));
    float *h_C_ref = (float*)malloc(M * N * sizeof(float));

    // 初始化矩阵A和B（填充随机值）
    for (int i = 0; i < M*K; ++i) h_A[i] = rand() % 10;
    for (int i = 0; i < K*N; ++i) h_B[i] = rand() % 10;

    // 计算CPU参考结果
    matrixMulCPU(h_A, h_B, h_C_ref);

    // 设备内存分配（转换为float4指针）
    float4 *d_A, *d_B;
    float  *d_C;
    cudaMalloc(&d_A, M * K * sizeof(float)); // 总字节数不变，但以float4形式访问
    cudaMalloc(&d_B, K * N * sizeof(float));
    cudaMalloc(&d_C, M * N * sizeof(float));

    // 数据拷贝到设备（直接拷贝整个float数组）
    cudaMemcpy(d_A, h_A, M * K * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, h_B, K * N * sizeof(float), cudaMemcpyHostToDevice);

    // 定义线程块和网格
    dim3 block(16, 16);
    dim3 grid((N + block.x - 1) / block.x, (M + block.y - 1) / block.y);

    // 启动核函数
    matrixMulGPU<<<grid, block>>>(d_A, d_B, d_C);

    // 拷贝结果回主机
    cudaMemcpy(h_C, d_C, M * N * sizeof(float), cudaMemcpyDeviceToHost);

    // 验证结果
    bool correct = true;
    for (int i = 0; i < M*N; ++i) {
        if (fabs(h_C[i] - h_C_ref[i]) > EPSILON) {
            printf("Mismatch at %d: CPU=%.2f, GPU=%.2f\n", i, h_C_ref[i], h_C[i]);
            correct = false;
            //break;
        }
    }
    printf("Result: %s\n", correct ? "PASS" : "FAIL");

    // 释放内存
    free(h_A); free(h_B); free(h_C); free(h_C_ref);
    cudaFree(d_A); cudaFree(d_B); cudaFree(d_C);

    return 0;
}